#!/usr/bin/env python3
"""
Stable Video Diffusion 环境配置脚本
自动检测和配置GPU环境，安装所需依赖
"""

import os
import sys
import subprocess
import platform
import logging
from pathlib import Path

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class EnvironmentSetup:
    def __init__(self):
        self.platform = platform.system()
        self.python_version = sys.version_info
        self.cuda_available = False
        self.cuda_version = None
        
    def check_python_version(self):
        """检查Python版本"""
        logger.info(f"检查Python版本: {sys.version}")
        if self.python_version.major == 3 and self.python_version.minor in [10, 11]:
            logger.info("✓ Python版本符合要求")
            return True
        else:
            logger.error("✗ 需要Python 3.10或3.11")
            return False
    
    def check_cuda(self):
        """检查CUDA环境"""
        logger.info("检查CUDA环境...")
        try:
            result = subprocess.run(
                ['nvidia-smi'], 
                capture_output=True, 
                text=True,
                shell=True
            )
            if result.returncode == 0:
                self.cuda_available = True
                logger.info("✓ NVIDIA GPU 检测成功")
                
                # 获取CUDA版本
                cuda_check = subprocess.run(
                    ['nvcc', '--version'],
                    capture_output=True,
                    text=True,
                    shell=True
                )
                if cuda_check.returncode == 0:
                    output = cuda_check.stdout
                    # 解析CUDA版本
                    for line in output.split('\n'):
                        if 'release' in line.lower():
                            self.cuda_version = line.split('release')[-1].strip().split(',')[0]
                            logger.info(f"✓ CUDA版本: {self.cuda_version}")
                            break
                return True
            else:
                logger.warning("✗ 未检测到NVIDIA GPU")
                return False
        except FileNotFoundError:
            logger.error("✗ nvidia-smi未找到，请安装NVIDIA驱动")
            return False
        except Exception as e:
            logger.error(f"✗ CUDA检查失败: {e}")
            return False
    
    def create_directories(self):
        """创建必要的目录结构"""
        logger.info("创建目录结构...")
        directories = [
            'models',
            'logs',
            'cache',
            'uploads',
            'outputs',
            'config'
        ]
        
        for dir_name in directories:
            dir_path = Path(dir_name)
            dir_path.mkdir(exist_ok=True)
            logger.info(f"  ✓ 创建目录: {dir_name}/")
    
    def install_dependencies(self):
        """安装Python依赖"""
        logger.info("安装Python依赖包...")
        
        # 创建requirements.txt
        requirements = """# Core dependencies
torch>=2.0.0
torchvision>=0.15.0
transformers>=4.35.0
diffusers>=0.24.0
accelerate>=0.25.0
safetensors>=0.4.0
omegaconf>=2.3.0

# Image and Video processing
opencv-python>=4.8.0
Pillow>=10.0.0
imageio>=2.31.0
imageio-ffmpeg>=0.4.9
moviepy>=1.0.3

# API Server
fastapi>=0.104.0
uvicorn>=0.24.0
python-multipart>=0.0.6
pydantic>=2.0.0
pydantic-settings>=2.0.0

# Utilities
numpy>=1.24.0
scipy>=1.11.0
tqdm>=4.65.0
pyyaml>=6.0
huggingface-hub>=0.19.0
einops>=0.7.0
xformers>=0.0.22  # 可选，用于内存优化

# Monitoring and Logging
prometheus-client>=0.19.0
loguru>=0.7.0

# Testing
pytest>=7.4.0
pytest-asyncio>=0.21.0
httpx>=0.25.0
"""
        
        req_file = Path('requirements.txt')
        req_file.write_text(requirements)
        logger.info("  ✓ 创建 requirements.txt")
        
        # 升级pip
        logger.info("  升级pip...")
        subprocess.run(['python', '-m', 'pip', 'install', '--upgrade', 'pip'], check=True)
        
        # 安装PyTorch (根据CUDA版本选择)
        if self.cuda_available:
            if self.cuda_version and self.cuda_version.startswith('11'):
                torch_cmd = f"python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118"
            elif self.cuda_version and self.cuda_version.startswith('12'):
                torch_cmd = f"python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121"
            else:
                torch_cmd = f"python -m pip install torch torchvision torchaudio"
            
            logger.info(f"  安装PyTorch (CUDA {self.cuda_version})...")
            subprocess.run(torch_cmd, shell=True, check=True)
        else:
            logger.info("  安装PyTorch (CPU版本)...")
            subprocess.run(f"python -m pip install torch torchvision torchaudio", shell=True, check=True)
        
        # 安装其他依赖
        logger.info("  安装其他依赖包...")
        subprocess.run(['python', '-m', 'pip', 'install', '-r', 'requirements.txt'], check=True)
        
        logger.info("✓ 所有依赖安装完成")
    
    def create_config_file(self):
        """创建默认配置文件"""
        logger.info("创建配置文件...")
        
        config_content = """# Stable Video Diffusion 配置文件

server:
  host: 0.0.0.0
  port: 8080
  workers: 2
  reload: false
  max_upload_size: 10485760  # 10MB
  
model:
  name: stabilityai/stable-video-diffusion-img2vid-xt
  variant: fp16
  device: cuda
  cache_dir: ./models
  use_safetensors: true
  enable_xformers: true  # 内存优化
  
inference:
  # 默认推理参数
  num_frames: 25  # 14 或 25
  decode_chunk_size: 8
  fps: 7
  motion_bucket_id: 127  # 0-255, 控制运动幅度
  noise_aug_strength: 0.02  # 0-1, 噪声增强
  min_guidance_scale: 1.0
  max_guidance_scale: 3.0
  seed: -1  # -1 表示随机
  
  # 性能设置
  batch_size: 1
  num_inference_steps: 25
  height: 576
  width: 1024
  
logging:
  level: INFO
  format: "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
  log_dir: ./logs
  max_file_size: 10485760  # 10MB
  backup_count: 5
  
cache:
  enable: true
  cache_dir: ./cache
  max_cache_size: 10737418240  # 10GB
  ttl: 3600  # 1小时
  
security:
  enable_auth: false
  api_key: ""
  max_requests_per_minute: 60
  allowed_origins:
    - "*"
  
monitoring:
  enable_metrics: true
  metrics_port: 9090
  
paths:
  upload_dir: ./uploads
  output_dir: ./outputs
  temp_dir: ./temp
"""
        
        config_file = Path('config/config.yaml')
        config_file.parent.mkdir(exist_ok=True)
        config_file.write_text(config_content)
        logger.info("  ✓ 创建 config/config.yaml")
    
    def verify_installation(self):
        """验证安装是否成功"""
        logger.info("\n验证安装...")
        
        try:
            # 测试PyTorch
            import torch
            logger.info(f"  ✓ PyTorch版本: {torch.__version__}")
            
            if torch.cuda.is_available():
                logger.info(f"  ✓ CUDA可用: {torch.cuda.get_device_name(0)}")
                logger.info(f"  ✓ CUDA设备数: {torch.cuda.device_count()}")
            else:
                logger.warning("  ⚠ CUDA不可用，将使用CPU模式")
            
            # 测试其他关键库
            import diffusers
            logger.info(f"  ✓ Diffusers版本: {diffusers.__version__}")
            
            import transformers
            logger.info(f"  ✓ Transformers版本: {transformers.__version__}")
            
            import fastapi
            logger.info(f"  ✓ FastAPI版本: {fastapi.__version__}")
            
            logger.info("\n✅ 环境配置成功！")
            return True
            
        except ImportError as e:
            logger.error(f"\n❌ 验证失败: {e}")
            return False
    
    def run(self):
        """运行完整的环境设置流程"""
        logger.info("="*50)
        logger.info("开始配置Stable Video Diffusion环境")
        logger.info("="*50)
        
        steps = [
            ("检查Python版本", self.check_python_version),
            ("检查CUDA环境", self.check_cuda),
            ("创建目录结构", self.create_directories),
            ("安装依赖包", self.install_dependencies),
            ("创建配置文件", self.create_config_file),
            ("验证安装", self.verify_installation)
        ]
        
        for step_name, step_func in steps:
            logger.info(f"\n步骤: {step_name}")
            logger.info("-" * 30)
            
            try:
                result = step_func()
                if result is False:
                    logger.warning(f"步骤 '{step_name}' 完成但有警告")
            except Exception as e:
                logger.error(f"步骤 '{step_name}' 失败: {e}")
                logger.info("\n环境配置未完成，请检查错误信息")
                return False
        
        logger.info("\n" + "="*50)
        logger.info("🎉 环境配置完成！")
        logger.info("="*50)
        logger.info("\n下一步:")
        logger.info("1. 运行 'python download_models.py' 下载模型")
        logger.info("2. 运行 'python svd_server.py' 启动服务")
        
        return True


if __name__ == "__main__":
    setup = EnvironmentSetup()
    success = setup.run()
    sys.exit(0 if success else 1)
